#pragma once

#include <Python.h>
#include <pcap/pcap.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netinet/ip.h>

/////////////////////////////////////////////////
/// \class SharedPyObjPtr
///
/// \brief Smart pointer for owner python object
///
/// Python use a reference count for each object,
/// the reference count of the object should decrease
/// one if owned the reference and don't use it anymore.
/////////////////////////////////////////////////
class SharedPyObjPtr
{
public:

	/////////////////////////////////////////////////
	/// \fn explicit SharedPyObjPtr(PyObject* pyObj = NULL)
	/// \param [in] pyObj python object with owner reference
	///
	/// Create a python object smart pointer with an optional python object pointer.
	/////////////////////////////////////////////////
	explicit SharedPyObjPtr(PyObject* pyObj = NULL)
		: m_pyObj(pyObj)
	{
	}

	explicit SharedPyObjPtr(const SharedPyObjPtr& otherPtr)
	{
		m_pyObj = otherPtr.m_pyObj;
		Py_XINCREF(m_pyObj);
	}

	SharedPyObjPtr& operator=(const SharedPyObjPtr& otherPtr)
	{
		Py_XDECREF(m_pyObj);
		m_pyObj = otherPtr.m_pyObj;
		Py_XINCREF(m_pyObj);
		return *this;
	}

	/// Get the original python object
	PyObject* get()
	{
		return m_pyObj;
	}

	/// Convertion for convert smart python pointer to original python pointer
	operator PyObject*()
	{
		return m_pyObj;
	}

	~SharedPyObjPtr()
	{
		Py_XDECREF(m_pyObj);
	}
private:
	PyObject* m_pyObj;	/// The original python object pointer
};

/////////////////////////////////////////////////
/// \class PyPacketFilter
///
/// \brief The python packet filter
///
/// Every packet pass into this class will analyze by
/// a python script.  The python script can be set in
/// runtime and will be applied immediately.
/////////////////////////////////////////////////
class PyPacketFilter
{
public:
	PyPacketFilter();
	virtual ~PyPacketFilter();

	/////////////////////////////////////////////////
	/// \fn bool SetPySourceFilter(const char* pyCode)
	/// \param pyCode python script to filter a packet
	/// \return true if the python source can be compiled with no error,
	///         false if the python source compiled with errors.
	///
	/// Set the python source code, it will be applied immediately.
	/// The script has several built-in variables:
	///   <i>packet</i>: the bytearray of the whole packet
	///   <i>localIP</i>: an integer represent the local IP address of
	///                   the client TCP connection in network byte order
	///   <i>localPort</i>: an integer represent the local port of
	///                     the client TCP connection in network byte order
	///   <i>peerIP</i>: an integer represent the client IP address of
	///                  the client TCP connection in network byte order
	///   <i>peerPort</i>: an integer represent the client port of
	///                    the client TCP connection in network byte order
	/// Note: the client TCP connection is used to set the python
	///       filter source and send the packet captured to the client.
	/// The script should set the boolean local variable <i>ret</i>
	///   to indicate whether the packet will be send to client or
	///   drop the packet.  <i>True</i> for send the packet to the
	///   packet processor, <i>False</i> for drop the packet.
	/////////////////////////////////////////////////
	bool SetPySourceFilter(const char* pyCode);

	/////////////////////////////////////////////////
	/// \fn bool operator()(const pcap_pkthdr *h, const u_char *bytes, sockaddr_in& sockaddr, sockaddr_in& peeraddr)
	/// \param pcap_pkthdr the packet information of a packet, declared in libpcap
	/// \param bytes the whole packet data
	/// \param sockaddr the local socket address of the client TCP connection
	/// \param peeraddr the remote peer address of the client TCP connection
	/// \return true for send the packet to the packet processor,
	///         false for drop the packet
	///
	/// Packet passed into this function will be determine
	/// whether the packet will be sent to the packet processor
	/// or drop it.
	/////////////////////////////////////////////////
	bool operator()(const pcap_pkthdr *h, const u_char *bytes, sockaddr_in& sockaddr, sockaddr_in& peeraddr);
private:
	SharedPyObjPtr m_pyFilter;	/// The python object of the compiled filter source
};
